# -*- coding: utf-8 -*-

import os
import cv2


IMAGE_HEIGHT = 256
IMAGE_WIDTH = 480
IMAGE_DEPTH = 3
GROUND_TRUTH_DEPTH = 1


def _ground_truth_process(ground_truth):
    _, ground_truth = cv2.threshold(ground_truth, 0, 255, cv2.THRESH_BINARY)
    ground_truth = cv2.medianBlur(ground_truth, ksize=7)
    ground_truth = cv2.dilate(ground_truth, kernel=(3, 3))
    return ground_truth


def generate_train_data():
    src_folder = "/home/lijun/DataBase/road_segment/origin"
    src_images_name = [name for name in os.listdir(src_folder)
                       if "jpg" in name and "_mask" not in name]
    src_ground_truth_files_name = [name[:name.rfind(".")] + "_mask.jpg"
                                   for name in src_images_name]
    dst_image_folder = "/home/lijun/DataBase/road_segment/480_256/images"
    dst_ground_truth_folder = "/home/lijun/DataBase/road_segment/480_256/ground_truth"
    index = 0
    for src_image_name, src_ground_truth_file_name in zip(src_images_name, src_ground_truth_files_name):
        index += 1
        src_image_path = os.path.join(src_folder, src_image_name)
        src_image = cv2.imread(src_image_path)
        dst_image = cv2.resize(src_image, dsize=(IMAGE_WIDTH, IMAGE_HEIGHT))
        dst_image_path = "{}/{:>07}.jpg".format(dst_image_folder, index)
        cv2.imwrite(dst_image_path, dst_image)

        src_ground_truth = cv2.imread(os.path.join(src_folder, src_ground_truth_file_name))
        dst_ground_truth = _ground_truth_process(src_ground_truth)
        dst_ground_truth = cv2.resize(dst_ground_truth, dsize=(IMAGE_WIDTH, IMAGE_HEIGHT))
        dst_ground_truth_path = "{}/{:>07}.jpg".format(dst_ground_truth_folder, index)
        cv2.imwrite(dst_ground_truth_path, dst_ground_truth)

        index += 1
        dst_image_path = "{}/{:>07}.jpg".format(dst_image_folder, index)
        dst_image = cv2.flip(dst_image, flipCode=1)  # horizontal flip
        cv2.imwrite(dst_image_path, dst_image)
        dst_ground_truth_path = "{}/{:>07}.jpg".format(dst_ground_truth_folder, index)
        dst_ground_truth = cv2.flip(dst_ground_truth, flipCode=1)
        cv2.imwrite(dst_ground_truth_path, dst_ground_truth)

        print(index)


def main():
    generate_train_data()


if __name__ == "__main__":
    main()
